#version 460

#extension GL_EXT_ray_tracing : require
#extension GL_EXT_ray_query : enable
#extension GL_GOOGLE_include_directive : require
#extension GL_EXT_nonuniform_qualifier : require
#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require
#extension GL_EXT_buffer_reference2 : require
#extension GL_EXT_scalar_block_layout : enable

#define RAY_TRACING
#include "../Brdf.glsl"

layout (set = 1, binding = 0) uniform accelerationStructureEXT topLevelAS;

#include "../RayQuery.glsl"
#define SOFT_SHADOWS
#define RAY_THROUGHPUT
#define SAMPLE_SKY_LIGHT
#include "../Lighting.glsl"

layout (location = 0) rayPayloadInEXT PathTracePayload hitPayload;

layout (location = 1) rayPayloadEXT PathTracePayload indirectPayload;

hitAttributeEXT vec2 hitAttribs;

layout (set = 1, binding = 1) uniform PerFrameUBO
{
    mat4 viewInverse;
    mat4 projInverse;
    mat4 projViewInverse;
    mat4 prevProjView;
    mat4 projView;
    vec4 cameraPosition;
    vec4 currentPrevJitter;
    Light light;
} perFrameUBO;

layout (set = 1, binding = 2) uniform samplerCube cubemapImage;

layout (buffer_reference, scalar) buffer Vertices {Vertex v[]; };

layout (buffer_reference, scalar) buffer Indices {ivec3 i[]; };

layout (set = 1, binding = 3, scalar) buffer ObjectDescBuffer { ObjectDesc data[]; } objectDescs;

layout (set = 1, binding = 4, scalar) buffer MaterialBuffer { Material data[]; } materials;

layout (push_constant) uniform PushConstants
{
    uint frame;
    uint maxRayBounces;
    float roughnessMultiplier;
} pushConstants;


vec3 IndirectLighting(vec3 Wo, vec3 N, vec3 P, vec3 F0, vec3 diffuse_color, float roughness, float metallic)
{
    vec3 Wi;
    float pdf;

    vec3 brdf = SampleUberBrdf(diffuse_color, F0, N, roughness, metallic, Wo, hitPayload.rng, Wi, pdf);

    float cos_theta = clamp(dot(N, Wi), 0.0, 1.0);

    indirectPayload.L = vec3(0.0f);
    indirectPayload.T = hitPayload.T * (brdf * cos_theta) / pdf;

    // Russian roulette
    float probability = max(indirectPayload.T.r, max(indirectPayload.T.g, indirectPayload.T.b));
    if (NextFloat(hitPayload.rng) > probability)
        return vec3(0.0f);

    // Add the energy we 'lose' by randomly terminating paths
    indirectPayload.T *= 1.0f / probability;
    indirectPayload.depth = hitPayload.depth + 1;
    indirectPayload.rng = hitPayload.rng;

    uint rayFlags = gl_RayFlagsOpaqueEXT;
    uint cullMask = 0xFF;
    float tmin = 0.0001;
    float tmax = 10000.0;
    vec3 origin = P.xyz;

    // Trace Ray
    traceRayEXT(topLevelAS,
                rayFlags,
                cullMask,
                0,
                0,
                0,
                origin,
                tmin,
                Wi,
                tmax,
                1);

    return indirectPayload.L;
}


void main()
{
    const ObjectDesc objectDesc = objectDescs.data[gl_InstanceCustomIndexEXT];
    const Material material = materials.data[gl_InstanceCustomIndexEXT];

    Indices indices = Indices(objectDesc.indexAddress);
    Vertices vertices = Vertices(objectDesc.vertexAddress);

    // Indices of the triangle
    ivec3 index = indices.i[gl_PrimitiveID];

    // Vertex of the triangle
    Vertex v0 = vertices.v[index.x];
    Vertex v1 = vertices.v[index.y];
    Vertex v2 = vertices.v[index.z];

    const vec3 barycentrics = vec3(1.0 - hitAttribs.x - hitAttribs.y, hitAttribs.x, hitAttribs.y);

    const vec3 position = vec3(gl_ObjectToWorldEXT * vec4(v0.position * barycentrics.x + v1.position * barycentrics.y + v2.position * barycentrics.z, 1.0));

    const vec3 N = normalize(vec3((v0.normal * barycentrics.x + v1.normal * barycentrics.y + v2.normal * barycentrics.z) * gl_WorldToObjectEXT));

    const vec3 Wo = -gl_WorldRayDirectionEXT;
    const vec3 R = reflect(-Wo, N);

    const vec3 albedo = material.albedo.rgb;
    const float roughness = material.roughnessMetallic.r;
    const float metallic = material.roughnessMetallic.g;

    const vec3 F0 = mix(vec3(0.04f), albedo, metallic);
    const vec3 diffuseColor = mix(albedo * (vec3(1.0f) - F0), vec3(0.0f), metallic);

    hitPayload.L += DirectLighting(perFrameUBO.light, Wo, N, position, F0, diffuseColor, roughness, hitPayload.T, NextVec2(hitPayload.rng), NextVec2(hitPayload.rng), cubemapImage);

    if ((hitPayload.depth + 1) < pushConstants.maxRayBounces)
        hitPayload.L += IndirectLighting(Wo, N, position, F0, diffuseColor, roughness, metallic);
}